/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2024-2024. All rights reserved.
 * Description:
 * Author: l30050806
 * Create: 2024-12-27
 */

#ifndef TRANSPORT_MEM_H
#define TRANSPORT_MEM_H

#include <hccl/hccl_types.h>
#include <atomic>
#include "dispatcher_pub.h"
#include "local_ipc_notify.h"
#include "notify_pool.h"
#include "hccl_socket.h"
#include "hccl_network.h"
#include "private_types.h"
#include "local_rdma_rma_buffer.h"
#include "local_ipc_rma_buffer.h"
#include "remote_rdma_rma_buffer.h"
#include "remote_ipc_rma_buffer.h"

namespace hccl {
struct RmaBufferSlice {
    std::shared_ptr<RmaBuffer> rmaBuffer{nullptr};
    void* addr{nullptr};
    u64 len{0};
    HcclMemType memType{HcclMemType::HCCL_MEM_TYPE_DEVICE};
};

class TransportMem {
public:
    explicit TransportMem(const std::unique_ptr<NotifyPool> &notifyPool, const HcclNetDevCtx &netDevCtx,
        const HcclDispatcher &dispatcher, u32 localRankId, u32 remoteRankId);
    virtual ~TransportMem();
    virtual HcclResult SetSocket(const std::shared_ptr<HcclSocket> &socket) = 0;
    virtual HcclResult Connect(s32 timeoutSec) = 0;
    virtual HcclResult Write(const RmaBufferSlice &remoteRmaBufferSlice, const RmaBufferSlice &localRmaBufferSlice, const rtStream_t &stream) = 0;
    virtual HcclResult Read(const RmaBufferSlice &localRmaBufferSlice, const RmaBufferSlice &remoteRmaBufferSlice, const rtStream_t &stream) = 0;
    virtual HcclResult AddOpFence(const rtStream_t &stream) = 0;

protected:
    const std::unique_ptr<NotifyPool> &notifyPool_;
    const hccl::NetDevContext *netDevCtx_{nullptr};
    DispatcherPub *dispatcher_{nullptr};
    u32 localRankId_{0};
    u32 remoteRankId_{0};
    std::shared_ptr<HcclSocket> socket_{nullptr};
};
}  // namespace hccl
#endif